from src.memory_handler import MemoryHandler
import pandas as pd

class SessionManager:
    def __init__(self):
        self.conversation_history = []
        self.memory_handler = MemoryHandler()  # 实例化 MemoryHandler

        # 加载之前保存的会话历史
        saved_history = self.memory_handler.get_memory("conversation_history")
        if saved_history:
            self.conversation_history = saved_history

    def add_to_history(self, user_input, model_response):
        """将用户输入和模型响应添加到对话历史中。"""
        # 如果响应是 DataFrame 或 Series，转换为可序列化的格式
        if isinstance(model_response, pd.DataFrame):
            model_response = model_response.to_dict(orient='list')  # 或者选择合适的序列化方式
        elif isinstance(model_response, pd.Series):
            model_response = model_response.to_dict()
        
        self.conversation_history.append({"user_input": user_input, "model_response": model_response})

    def get_conversation_history(self):
        """返回当前的对话历史。"""
        return self.conversation_history

    def clear_history(self):
        """清除对话历史。"""
        self.conversation_history = []
        self.memory_handler.update_memory("conversation_history", self.conversation_history)

    def format_for_model(self):
        """将对话历史格式化为字符串，以传递给模型。"""
        formatted_history = ""
        for entry in self.conversation_history:
            formatted_history += f"User: {entry['user_input']}\n"
            formatted_history += f"Model: {entry['model_response']}\n"
        return formatted_history

    def save_memory(self):
        """保存当前对话历史到记忆中。"""
        self.memory_handler.update_memory("conversation_history", self.conversation_history)

